import cv2
import numpy as np
from pyimagesearch import transform
class ImageProcessor:
    def __init__(self, target_size=(1654, 2339)):  # A4纸比例
        self.target_size = target_size
        self.MORPH = 9
        self.CANNY = 84
        self.MIN_QUAD_AREA_RATIO = 0.25
        self.MAX_QUAD_ANGLE_RANGE = 40
        
    def detect_document(self, image):
        """检测文档边缘"""
        # 图像预处理
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        gray = cv2.GaussianBlur(gray, (7,7), 0)
        
        # 形态学处理
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (self.MORPH, self.MORPH))
        dilated = cv2.morphologyEx(gray, cv2.MORPH_CLOSE, kernel)
        
        # Canny边缘检测
        edged = cv2.Canny(dilated, 0, self.CANNY)
        
        # 轮廓检测
        contours, _ = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 按面积排序，取最大的轮廓
        if not contours:
            return None
            
        contours = sorted(contours, key=cv2.contourArea, reverse=True)
        doc_contour = None
        
        for contour in contours:
            # 多边形近似
            approx = cv2.approxPolyDP(contour, 0.02 * cv2.arcLength(contour, True), True)
            if len(approx) == 4:
                doc_contour = approx
                break
                
        return doc_contour

    def detect_text_lines(self, image):
        """检测文本行"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 自适应二值化
        binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                     cv2.THRESH_BINARY_INV, 25, 15)
        
        # 形态学操作突出文本行
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 1))
        text_lines = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        
        return text_lines, binary

    def perspective_transform(self, image, corners):
        """透视变换矫正"""
        try:
            if corners is None or len(corners) != 4:
                return image
                
            # 获取图像尺寸
            h, w = image.shape[:2]
            
            # 排序角点
            rect = transform.order_points(corners.astype("float32"))
            (tl, tr, br, bl) = rect
            
            # 计算新尺寸
            width_a = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
            width_b = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
            max_width = max(int(width_a), int(width_b))
            
            height_a = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
            height_b = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
            max_height = max(int(height_a), int(height_b))
            
            # 添加最小尺寸检查
            min_size = 100
            if max_width < min_size or max_height < min_size:
                print("检测到的文档区域过小")
                return image
                
            # 添加最大尺寸限制
            max_size = 4000
            if max_width > max_size or max_height > max_size:
                scale = max_size / max(max_width, max_height)
                max_width = int(max_width * scale)
                max_height = int(max_height * scale)
            
            # 保持A4纸比例
            target_ratio = 1.414  # A4纸比例
            current_ratio = max_width / max_height
            
            if current_ratio > target_ratio:
                max_height = int(max_width / target_ratio)
            else:
                max_width = int(max_height * target_ratio)
            
            # 目标坐标
            dst = np.array([
                [0, 0],
                [max_width - 1, 0],
                [max_width - 1, max_height - 1],
                [0, max_height - 1]
            ], dtype="float32")
            
            M = cv2.getPerspectiveTransform(rect, dst)
            warped = cv2.warpPerspective(image, M, (max_width, max_height))
            
            return warped
            
        except Exception as e:
            print(f"透视变换失败: {str(e)}")
            return image
            
    def _basic_perspective_transform(self, image, corners):
        """基础透视变换（作为备选方案）"""
        # 原来的透视变换代码作为备选
        rect = transform.order_points(corners.astype("float32"))
        (tl, tr, br, bl) = rect
        
        width_a = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
        width_b = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
        max_width = max(int(width_a), int(width_b))
        
        height_a = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
        height_b = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
        max_height = max(int(height_a), int(height_b))
        
        dst = np.array([
            [0, 0],
            [max_width - 1, 0],
            [max_width - 1, max_height - 1],
            [0, max_height - 1]
        ], dtype="float32")
        
        M = cv2.getPerspectiveTransform(rect, dst)
        return cv2.warpPerspective(image, M, (max_width, max_height))

    def enhance_image(self, image):
        """图像增强处理"""
        try:
            # 检查图像是否为空
            if image is None or image.size == 0:
                return None
                
            # 转灰度图
            if len(image.shape) == 3:
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:
                gray = image
                
            # 对比度增强
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
            gray = clahe.apply(gray)
            
            # USM锐化
            blurred = cv2.GaussianBlur(gray, (0, 0), 3)
            sharpened = cv2.addWeighted(gray, 1.5, blurred, -0.5, 0)
            
            # 自适应二值化参数优化
            binary = cv2.adaptiveThreshold(
                sharpened, 
                255, 
                cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                cv2.THRESH_BINARY, 
                25,  # 增大块大小
                12   # 调整常数
            )
            
            # 降噪
            kernel = np.ones((2,2), np.uint8)
            binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
            
            return binary
            
        except Exception as e:
            print(f"图像增强失败: {str(e)}")
            return image
        
    def resize_to_target(self, image):
        """调整图像大小到目标尺寸，保持宽高比"""
        try:
            if image is None:
                return None
                
            # 计算目标宽高比
            target_ratio = self.target_size[0] / self.target_size[1]
            # 计算原图宽高比
            h, w = image.shape[:2]
            image_ratio = w / h
            
            # 根据宽高比决定缩放方式
            if image_ratio > target_ratio:
                # 图片较宽，以高度为准
                new_width = int(self.target_size[1] * image_ratio)
                resized = cv2.resize(image, (new_width, self.target_size[1]), 
                                   interpolation=cv2.INTER_AREA)
                # 居中裁剪
                start = (new_width - self.target_size[0]) // 2
                resized = resized[:, start:start + self.target_size[0]]
            else:
                # 图片较高，以宽度为准
                new_height = int(self.target_size[0] / image_ratio)
                resized = cv2.resize(image, (self.target_size[0], new_height), 
                                   interpolation=cv2.INTER_AREA)
                # 居中裁剪
                start = (new_height - self.target_size[1]) // 2
                resized = resized[start:start + self.target_size[1], :]
                
            return resized
            
        except Exception as e:
            print(f"调整图像大小失败: {str(e)}")
            return image
        
    def process_image(self, image):
        """完整的图像处理流程"""
        try:
            if image is None or image.size == 0:
                return None
                
            # 1. 检测文档边缘
            corners = self.detect_document(image)
            if corners is None:
                # 如果检测失败，直接处理原图
                enhanced = self.enhance_image(image)
                # return self.resize_to_target(enhanced)
                return enhanced
                
            # 2. 透视变换
            warped = self.perspective_transform(image, corners.reshape(4, 2))
            
            # 3. 图像增强
            enhanced = self.enhance_image(warped)
            if enhanced is None:
                return None
                
            # 4. 调整到统一大小
            # final = self.resize_to_target(enhanced)
            
            return enhanced
            
        except Exception as e:
            print(f"图像处理失败: {str(e)}")
            return None